{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a Custom Operator in C++\n",
    "\n",
    "DALI allows you to create a custom operator in C++ and load it at runtime. Here are several reasons you might need to write your custom operator:\n",
    "\n",
    "- DALI does not support the operation that you want to perform and it cannot be expressed by a composition of other operators.\n",
    "- You want to write an operator that depends on a third-party library.\n",
    "- You want to optimize your pipeline by providing a manually fused operation in C++.\n",
    "\n",
    "In this tutorial, we will walk you through the process of writing, compiling, and loading a plugin with a DALI custom operator. For demonstration purposes we will provide a CPU and a GPU implementation for the `CustomDummy` operator. The implementation only copies the input data to the output without any modifications.\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "- DALI is installed from the binary distribution or compiled the from source.\n",
    "- You can write in C++.\n",
    "- You have a basic knowledge of CMake."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Operator Definition\n",
    "\n",
    "1. Declare the operator in a header file.\n",
    "\n",
    "2. Provide common Setup functions.\n",
    "\n",
    "The implementation `SetupImpl` can be shared across backends. `SetupImpl` provides the shape and type description of the output based on the input; the return value determines whether the executor should allocate the storage for the operator's output before calling `RunImpl`. `HasContiguousOutputs` declares that the outputs of the operator are contiguous - this is usually true, except for operators which shuffle the batch without copying or legacy operators utilizing `SampleWorkspace` - and that's why we can rely on the default implementation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#ifndef EXAMPLE_DUMMY_H_\n",
      "#define EXAMPLE_DUMMY_H_\n",
      "\n",
      "#include <vector>\n",
      "\n",
      "#include \"dali/pipeline/operator/operator.h\"\n",
      "\n",
      "namespace other_ns {\n",
      "\n",
      "template <typename Backend>\n",
      "class Dummy : public ::dali::Operator<Backend> {\n",
      " public:\n",
      "  inline explicit Dummy(const ::dali::OpSpec &spec) :\n",
      "    ::dali::Operator<Backend>(spec) {}\n",
      "\n",
      "  virtual inline ~Dummy() = default;\n",
      "\n",
      "  Dummy(const Dummy&) = delete;\n",
      "  Dummy& operator=(const Dummy&) = delete;\n",
      "  Dummy(Dummy&&) = delete;\n",
      "  Dummy& operator=(Dummy&&) = delete;\n",
      "\n",
      " protected:\n",
      "  bool SetupImpl(std::vector<::dali::OutputDesc> &output_desc,\n",
      "                 const ::dali::Workspace &ws) override {\n",
      "    const auto &input = ws.Input<Backend>(0);\n",
      "    output_desc.resize(1);\n",
      "    output_desc[0] = {input.shape(), input.type()};\n",
      "    return true;\n",
      "  }\n",
      "\n",
      "  void RunImpl(::dali::Workspace &ws) override;\n",
      "};\n",
      "\n",
      "}  // namespace other_ns\n",
      "\n",
      "#endif  // EXAMPLE_DUMMY_H_\n"
     ]
    }
   ],
   "source": [
    "! cat customdummy/dummy.h"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CPU Operator Implementation\n",
    "\n",
    "1. Provide the CPU implementation in a C++ implementation file by overriding the RunImpl method for Workspace.\n",
    "\n",
    "2. Register the schema for the custom operator with DALI_SCHEMA macro and register the CPU version of the operator with DALI_REGISTER_OPERATOR.\n",
    "\n",
    "In RunImpl we obtain access to the entire batch that is processed. We get the reference to the CPU thread pool from the workspace `ws` and create tasks that will copy samples from input to output in parallel. The tasks will be ordered by the thread pool from the longest to the shortest, based on the tensor size, to best utilize the worker threads.\n",
    "\n",
    "The outputs are already allocated as we provided the SetupImpl function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#include \"dummy.h\"\n",
      "\n",
      "namespace other_ns {\n",
      "\n",
      "template <>\n",
      "void Dummy<::dali::CPUBackend>::RunImpl(::dali::Workspace &ws) {\n",
      "  const auto &input = ws.Input<::dali::CPUBackend>(0);\n",
      "  auto &output = ws.Output<::dali::CPUBackend>(0);\n",
      "\n",
      "  ::dali::TypeInfo type = input.type_info();\n",
      "  auto &tp = ws.GetThreadPool();\n",
      "  const auto &in_shape = input.shape();\n",
      "  for (int sample_id = 0; sample_id < in_shape.num_samples(); sample_id++) {\n",
      "    tp.AddWork(\n",
      "        [&, sample_id](int thread_id) {\n",
      "          type.Copy<::dali::CPUBackend, ::dali::CPUBackend>(\n",
      "                            output.raw_mutable_tensor(sample_id),\n",
      "                            input.raw_tensor(sample_id),\n",
      "                            in_shape.tensor_size(sample_id), 0);\n",
      "        },\n",
      "        in_shape.tensor_size(sample_id));\n",
      "  }\n",
      "  tp.RunAll();\n",
      "}\n",
      "\n",
      "}  // namespace other_ns\n",
      "\n",
      "DALI_REGISTER_OPERATOR(CustomDummy, ::other_ns::Dummy<::dali::CPUBackend>,\n",
      "                       ::dali::CPU);\n",
      "\n",
      "DALI_SCHEMA(CustomDummy)\n",
      "    .DocStr(\"Make a copy of the input tensor\")\n",
      "    .NumInput(1)\n",
      "    .NumOutput(1);\n"
     ]
    }
   ],
   "source": [
    "! cat customdummy/dummy.cc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPU operator implementation\n",
    "\n",
    "1. Provide a GPU implementation in a CUDA implementation file by overriding the RunImpl method for Workspace.\n",
    "\n",
    "2. Register the GPU version of the operator with DALI_REGISTER_OPERATOR macro.\n",
    "\n",
    "As it was the case for the CPU implementation, we obtain the entire batch in the RunImpl function. The outputs are already allocated based on the return value of SetupImpl function that was provided earlier.\n",
    "\n",
    "It is important that we issue the GPU operations on the stream provided by the workspace. Here we copy the batch using cudaMemcpyAsync."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#include <cuda_runtime_api.h>\n",
      "#include \"dummy.h\"\n",
      "\n",
      "namespace other_ns {\n",
      "\n",
      "template<>\n",
      "void Dummy<::dali::GPUBackend>::RunImpl(::dali::Workspace &ws) {\n",
      "  const auto &input = ws.Input<::dali::GPUBackend>(0);\n",
      "  const auto &shape = input.shape();\n",
      "  auto &output = ws.Output<::dali::GPUBackend>(0);\n",
      "  for (int sample_idx = 0; sample_idx < shape.num_samples(); sample_idx++) {\n",
      "    CUDA_CALL(cudaMemcpyAsync(\n",
      "            output.raw_mutable_tensor(sample_idx),\n",
      "            input.raw_tensor(sample_idx),\n",
      "            shape[sample_idx].num_elements() * input.type_info().size(),\n",
      "            cudaMemcpyDeviceToDevice,\n",
      "            ws.stream()));\n",
      "  }\n",
      "}\n",
      "\n",
      "}  // namespace other_ns\n",
      "\n",
      "DALI_REGISTER_OPERATOR(CustomDummy, ::other_ns::Dummy<::dali::GPUBackend>,\n",
      "                       ::dali::GPU);\n"
     ]
    }
   ],
   "source": [
    "! cat customdummy/dummy.cu"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building the Plugin\n",
    "1. Specify the build configuration.\n",
    "\n",
    "To retrieve the build configuration parameters use nvidia.dali.sysconfig."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nvidia.dali.sysconfig as sysconfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/nvidia/dali/include\n"
     ]
    }
   ],
   "source": [
    "print(sysconfig.get_include_dir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.8/dist-packages/nvidia/dali\n"
     ]
    }
   ],
   "source": [
    "print(sysconfig.get_lib_dir())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['-I/usr/local/lib/python3.8/dist-packages/nvidia/dali/include', '-D_GLIBCXX_USE_CXX11_ABI=1']\n"
     ]
    }
   ],
   "source": [
    "print(sysconfig.get_compile_flags())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['-L/usr/local/lib/python3.8/dist-packages/nvidia/dali', '-ldali']\n"
     ]
    }
   ],
   "source": [
    "print(sysconfig.get_link_flags())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "**Important:** *Only one version of libdali.so should be loaded in the process at the same time. A plugin must be linked against the exact same library in the DALI's Python package directory that you intend to use to load your plugin. As a result of this limitation, when you upgrade your DALI version you must link your plugin against the new library again.*\n",
    "\n",
    "2. In this example, we used CMake to build the plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cmake_minimum_required(VERSION 3.25.2)\n",
      "set(CMAKE_CUDA_ARCHITECTURES \"50;60;70;80;90\")\n",
      "\n",
      "project(custom_dummy_plugin LANGUAGES CUDA CXX C)\n",
      "\n",
      "set(CMAKE_CXX_STANDARD 20)\n",
      "set(CMAKE_CXX_STANDARD_REQUIRED ON)\n",
      "set(CMAKE_CXX_EXTENSIONS OFF)\n",
      "set(CMAKE_C_STANDARD 11)\n",
      "\n",
      "# TODO(klecki): When the test container gets a CMake that supports C++20 as a proper option,\n",
      "# swap those lines\n",
      "# set(CMAKE_CUDA_STANDARD 20)\n",
      "# set(CMAKE_CUDA_STANDARD_REQUIRED ON)\n",
      "set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} -std=c++20\")\n",
      "\n",
      "include_directories(SYSTEM \"${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}\")\n",
      "\n",
      "execute_process(\n",
      "        COMMAND python -c \"import nvidia.dali as dali; \\\n",
      "                           print(dali.sysconfig.get_lib_dir())\"\n",
      "        OUTPUT_VARIABLE DALI_LIB_DIR)\n",
      "string(STRIP ${DALI_LIB_DIR} DALI_LIB_DIR)\n",
      "\n",
      "execute_process(\n",
      "        COMMAND python -c \"import nvidia.dali as dali; print(\\\" \\\n",
      "                          \\\".join(dali.sysconfig.get_compile_flags()))\"\n",
      "        OUTPUT_VARIABLE DALI_COMPILE_FLAGS)\n",
      "string(STRIP ${DALI_COMPILE_FLAGS} DALI_COMPILE_FLAGS)\n",
      "\n",
      "set(CMAKE_CXX_FLAGS \"${CMAKE_CXX_FLAGS} ${DALI_COMPILE_FLAGS} \")\n",
      "set(CMAKE_CUDA_FLAGS \"${CMAKE_CUDA_FLAGS} ${DALI_COMPILE_FLAGS} \")\n",
      "link_directories(\"${DALI_LIB_DIR}\")\n",
      "\n",
      "add_library(dali_customdummy SHARED dummy.cc dummy.cu)\n",
      "target_link_libraries(dali_customdummy dali)\n"
     ]
    }
   ],
   "source": [
    "! cat customdummy/CMakeLists.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. We are now ready to compile the plugin that contains the `CustomDummy` custom operator."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-- The CUDA compiler identification is NVIDIA 12.4.131\n",
      "-- The CXX compiler identification is GNU 9.4.0\n",
      "-- The C compiler identification is GNU 9.4.0\n",
      "-- Detecting CUDA compiler ABI info\n",
      "-- Detecting CUDA compiler ABI info - done\n",
      "-- Check for working CUDA compiler: /usr/local/cuda/bin/nvcc - skipped\n",
      "-- Detecting CUDA compile features\n",
      "-- Detecting CUDA compile features - done\n",
      "-- Detecting CXX compiler ABI info\n",
      "-- Detecting CXX compiler ABI info - done\n",
      "-- Check for working CXX compiler: /usr/bin/c++ - skipped\n",
      "-- Detecting CXX compile features\n",
      "-- Detecting CXX compile features - done\n",
      "-- Detecting C compiler ABI info\n",
      "-- Detecting C compiler ABI info - done\n",
      "-- Check for working C compiler: /usr/bin/cc - skipped\n",
      "-- Detecting C compile features\n",
      "-- Detecting C compile features - done\n",
      "-- Configuring done (5.2s)\n",
      "-- Generating done (0.0s)\n",
      "-- Build files have been written to: /dali/docs/examples/custom_operations/custom_operator/customdummy/build\n",
      "[ 33%] \u001b[32mBuilding CXX object CMakeFiles/dali_customdummy.dir/dummy.cc.o\u001b[0m\n",
      "[ 66%] \u001b[32mBuilding CUDA object CMakeFiles/dali_customdummy.dir/dummy.cu.o\u001b[0m\n",
      "[100%] \u001b[32m\u001b[1mLinking CXX shared library libdali_customdummy.so\u001b[0m\n",
      "[100%] Built target dali_customdummy\n"
     ]
    }
   ],
   "source": [
    "! rm -rf customdummy/build\n",
    "! mkdir -p customdummy/build\n",
    "# CLICOLOR_FORCE=1 is set by default by Jupyter, but can cause problems with CMake\n",
    "! cd customdummy/build && \\\n",
    "  CLICOLOR_FORCE=0 cmake .. && \\\n",
    "  make"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "4. After the build is complete we have a dynamic library file that is ready to use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "customdummy/build/libdali_customdummy.so\n"
     ]
    }
   ],
   "source": [
    "! ls customdummy/build/*.so"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Importing the Plugin\n",
    "\n",
    "1. We can see that there is no such operator called `custom_dummy`.\n",
    "\n",
    "**Note**: Operations available in `nvidia.dali.fn` are automatically converted from camel case to snake case, while the legacy operator objects in `nvidia.dali.ops` keep the camel case format (Example: `fn.custom_dummy` vs. `ops.CustomDummy`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Error: module 'nvidia.dali.fn' has no attribute 'custom_dummy'\n"
     ]
    }
   ],
   "source": [
    "import nvidia.dali.fn as fn\n",
    "\n",
    "try:\n",
    "    help(fn.custom_dummy)\n",
    "except Exception as e:\n",
    "    print(\"Error: \" + str(e))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "2. Load the plugin."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "import nvidia.dali.plugin_manager as plugin_manager\n",
    "\n",
    "plugin_manager.load_library(\"./customdummy/build/libdali_customdummy.so\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "3. Verify that the new operator is available."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Help on function custom_dummy in module nvidia.dali.fn:\n",
      "\n",
      "custom_dummy(__input, /, *, bytes_per_sample_hint=[0], preserve=False, seed=-1, device=None, name=None)\n",
      "    Make a copy of the input tensor\n",
      "    \n",
      "    Supported backends\n",
      "     * 'cpu'\n",
      "     * 'gpu'\n",
      "    \n",
      "    \n",
      "    Args\n",
      "    ----\n",
      "    __input : TensorList\n",
      "        Input to the operator.\n",
      "    \n",
      "    \n",
      "    Keyword args\n",
      "    ------------\n",
      "    bytes_per_sample_hint : int or list of int, optional, default = `[0]`\n",
      "        Output size hint, in bytes per sample.\n",
      "        \n",
      "        If specified, the operator's outputs residing in GPU or page-locked host memory will be preallocated\n",
      "        to accommodate a batch of samples of this size.\n",
      "    preserve : bool, optional, default = `False`\n",
      "        Prevents the operator from being removed from the\n",
      "        graph even if its outputs are not used.\n",
      "    seed : int, optional, default = `-1`\n",
      "        Random seed.\n",
      "        \n",
      "        If not provided, it will be populated based on the global seed of the pipeline.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "help(fn.custom_dummy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the sake of completeness, it is worth mentioning that even if it discouraged, it is also possible to access the custom operator through the legacy operator object API (nvidia.dali.ops.CustomDummy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preloading the plugins\n",
    "\n",
    "DALI can preload plugins automatically via the environment variable DALI_PRELOAD_PLUGINS. If provided, the variable is interpreted as a list of paths separated by a colon. Each of the elements in the list can be either a directory or a library path. Directories are searched for files matching the libdali_*.so pattern. If not set, the \"default\" plugin directory is scanned, which is a \"plugin\" directory inside the DALI library installation directory."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
